import numpy as np
import pygame
from Box2D import b2_pi


def manhattan(x, y):
    return np.sum(np.abs(x - y))


def drawTheWorld(worldContext):
    worldContext['screen'].fill((0, 0, 0, 0))
    bodiesInTheWorld = worldContext['bodiesInTheWorld']
    for bodyKey in bodiesInTheWorld.keys():
        body = bodiesInTheWorld[bodyKey]
        for fixture in body.fixtures:
            shape = fixture.shape
            vertices = [(body.transform * v) * worldContext['PPM'] for v in shape.vertices]
            vertices = [(v[0], worldContext['SCREEN_HEIGHT'] - v[1]) for v in vertices]
            pygame.draw.polygon(worldContext['screen'], worldContext['colors'][body.type], vertices)
    worldContext['world'].Step(worldContext['TIME_STEP'], 10, 10)
    pygame.display.flip()
    worldContext['clock'].tick(worldContext['TARGET_FPS'])
    pygame.image.save(worldContext['screen'], 'outputImages/temp.png')


def changeTheWorldByAction(action, worldContext):
    bodiesInTheWorld = worldContext['bodiesInTheWorld']
    dynamic_body = bodiesInTheWorld['dynamic_body']
    if action == 0:
        dynamic_body.position -= (0.1, 0)
    elif action == 1:
        dynamic_body.position += (0.1, 0)
    elif action == 2:
        dynamic_body.position += (0, 0.1)
    elif action == 3:
        dynamic_body.position -= (0, 0.1)
    elif action == 4:
        dynamic_body.angle += 0.1 * b2_pi
    elif action == 5:
        dynamic_body.angle -= 0.1 * b2_pi


def encodeRGB(rgb):
    result = rgb[0] + rgb[1] * 255 + rgb[2] * 255 * 255
    return result


def encodeCoord(coord, width):
    h, w, _ = coord
    order = h * width + w
    return order


def encodeImage(img):
    h, w, _ = img.shape
    result = np.zeros((h, w, 1), dtype='int32')
    for i in range(h):
        for j in range(w):
            result[i, j] = encodeRGB(img[i, j])
    return result
